{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79dfbe1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | default_exp _components.aiokafka_consumer_loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aaf843a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/kumaran/.local/lib/python3.11/site-packages/pydantic/_internal/_config.py:257: UserWarning: Valid config keys have changed in V2:\n",
      "* 'json_encoders' has been removed\n",
      "  warnings.warn(message, UserWarning)\n"
     ]
    }
   ],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "from asyncio import iscoroutinefunction, Task  # do not use the version from inspect\n",
    "from typing import *\n",
    "from dataclasses import dataclass\n",
    "\n",
    "import asyncer\n",
    "from aiokafka.structs import ConsumerRecord\n",
    "from pydantic import BaseModel\n",
    "\n",
    "import fastkafka._aiokafka_imports\n",
    "from fastkafka._components.logger import get_logger\n",
    "from fastkafka._components.meta import delegates, export\n",
    "from fastkafka._components.task_streaming import get_executor, StreamExecutor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a446cf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import pytest\n",
    "from datetime import datetime, timedelta\n",
    "from unittest.mock import AsyncMock, MagicMock, Mock, call, patch, create_autospec\n",
    "\n",
    "import anyio\n",
    "from aiokafka.structs import TopicPartition\n",
    "from pydantic import Field, HttpUrl, NonNegativeInt\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from fastkafka._components.helpers import true_after\n",
    "from fastkafka._components.logger import suppress_timestamps\n",
    "from fastkafka._helpers import produce_messages\n",
    "from fastkafka.encoder import avro_decoder, avro_encoder, json_decoder\n",
    "from fastkafka.testing import ApacheKafkaBroker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0e9e4c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | notest\n",
    "# allows async calls in notebooks\n",
    "\n",
    "import nest_asyncio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53542175",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | notest\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af85a823",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "logger = get_logger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92feb585",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: ok\n"
     ]
    }
   ],
   "source": [
    "suppress_timestamps()\n",
    "logger = get_logger(__name__, level=20)\n",
    "logger.info(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeaf0d5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyMessage(BaseModel):\n",
    "    url: HttpUrl = Field(..., example=\"http://www.acme.com\", description=\"Url example\")\n",
    "    port: NonNegativeInt = Field(1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68cf645d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@dataclass\n",
    "@export(\"fastkafka\")\n",
    "class EventMetadata:\n",
    "    \"\"\"A class for encapsulating Kafka record metadata.\n",
    "\n",
    "    Args:\n",
    "        topic: The topic this record is received from\n",
    "        partition: The partition from which this record is received\n",
    "        offset: The position of this record in the corresponding Kafka partition\n",
    "        timestamp: The timestamp of this record\n",
    "        timestamp_type: The timestamp type of this record\n",
    "        key: The key (or `None` if no key is specified)\n",
    "        value: The value\n",
    "        serialized_key_size: The size of the serialized, uncompressed key in bytes\n",
    "        serialized_value_size: The size of the serialized, uncompressed value in bytes\n",
    "        headers: The headers\n",
    "    \"\"\"\n",
    "\n",
    "    topic: str\n",
    "    partition: int\n",
    "    offset: int\n",
    "    timestamp: int\n",
    "    timestamp_type: int\n",
    "    key: Optional[bytes]\n",
    "    value: Optional[bytes]\n",
    "    checksum: int\n",
    "    serialized_key_size: int\n",
    "    serialized_value_size: int\n",
    "    headers: Sequence[Tuple[str, bytes]]\n",
    "\n",
    "    @staticmethod\n",
    "    def create_event_metadata(record: ConsumerRecord) -> \"EventMetadata\":  # type: ignore\n",
    "        \"\"\"Creates an instance of EventMetadata from a ConsumerRecord.\n",
    "\n",
    "        Args:\n",
    "            record: The Kafka ConsumerRecord.\n",
    "\n",
    "        Returns:\n",
    "            The created EventMetadata instance.\n",
    "        \"\"\"\n",
    "        return EventMetadata(\n",
    "            topic=record.topic,\n",
    "            partition=record.partition,\n",
    "            offset=record.offset,\n",
    "            timestamp=record.timestamp,\n",
    "            timestamp_type=record.timestamp_type,\n",
    "            value=record.value,\n",
    "            checksum=record.checksum,\n",
    "            key=record.key,\n",
    "            serialized_key_size=record.serialized_key_size,\n",
    "            serialized_value_size=record.serialized_value_size,\n",
    "            headers=record.headers,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4e0d145",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_consumer_record(topic: str, partition: int, msg: BaseModel):\n",
    "    record = ConsumerRecord(\n",
    "        topic=topic,\n",
    "        partition=partition,\n",
    "        offset=0,\n",
    "        timestamp=0,\n",
    "        timestamp_type=0,\n",
    "        key=None,\n",
    "        value=msg.model_dump_json().encode(\"utf-8\")\n",
    "        if hasattr(msg, \"json\")\n",
    "        else msg.encode(\"utf-8\"),\n",
    "        checksum=0,\n",
    "        serialized_key_size=0,\n",
    "        serialized_value_size=0,\n",
    "        headers=[],\n",
    "    )\n",
    "    return record"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bf5ffd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "meta = EventMetadata.create_event_metadata(create_consumer_record(\"topic\", 1, MyMessage(url=\"http://www.acme.com\", port=22)))\n",
    "assert meta.topic == \"topic\"\n",
    "assert meta.partition == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bb76fa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "AsyncConsume = Callable[[Union[List[BaseModel], BaseModel]], Awaitable[None]]\n",
    "AsyncConsumeMeta =  Callable[[Union[List[BaseModel], BaseModel], Union[List[EventMetadata], EventMetadata]], Awaitable[None]]\n",
    "SyncConsume = Callable[[Union[List[BaseModel], BaseModel]], None]\n",
    "SyncConsumeMeta =  Callable[[Union[List[BaseModel], BaseModel], Union[List[EventMetadata], EventMetadata]], None]\n",
    "\n",
    "ConsumeCallable = Union[\n",
    "    AsyncConsume, AsyncConsumeMeta, SyncConsume, SyncConsumeMeta\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01f24d56",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _callback_parameters_wrapper(\n",
    "    callback: Union[AsyncConsume, AsyncConsumeMeta]\n",
    ") -> AsyncConsumeMeta:\n",
    "    \"\"\"Wraps an async callback and filters the arguments to pass based on if the function accepts EventMetadata as argument\n",
    "\n",
    "    Args:\n",
    "        callback: async callable that will be wrapped\n",
    "\n",
    "    Returns:\n",
    "        Wrapped callback with filtered params\n",
    "    \"\"\"\n",
    "\n",
    "    async def _params_wrap(\n",
    "        msg: Union[BaseModel, List[BaseModel]],\n",
    "        meta: Union[EventMetadata, List[EventMetadata]],\n",
    "        callback: Union[AsyncConsume, AsyncConsumeMeta] = callback,\n",
    "    ) -> None:\n",
    "        types = list(get_type_hints(callback).values())\n",
    "        args: List[Union[BaseModel, List[BaseModel], EventMetadata, List[EventMetadata]]] = [msg]\n",
    "        if EventMetadata in types:\n",
    "            args.insert(types.index(EventMetadata), meta)\n",
    "        if List[EventMetadata] in types:\n",
    "            args.insert(types.index(List[EventMetadata]), meta)\n",
    "        await callback(*args)  # type: ignore\n",
    "\n",
    "    return _params_wrap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b833bd5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def without_meta(msg: BaseModel):\n",
    "    assert msg == \"Example_msg\"\n",
    "\n",
    "with pytest.raises(TypeError) as e:\n",
    "    await without_meta(\"Example_msg\", \"Some_meta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ba7b7f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "@_callback_parameters_wrapper\n",
    "async def without_meta(msg: BaseModel):\n",
    "    assert msg == \"Example_msg\"\n",
    "\n",
    "await without_meta(\"Example_msg\", \"Some_meta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "859b390d",
   "metadata": {},
   "outputs": [],
   "source": [
    "@_callback_parameters_wrapper\n",
    "async def with_meta(msg: BaseModel, meta: EventMetadata):\n",
    "    assert msg == \"Example_msg\"\n",
    "    assert meta == \"Some_meta\"\n",
    "\n",
    "await with_meta(\"Example_msg\", \"Some_meta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec76eaf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "@_callback_parameters_wrapper\n",
    "async def with_meta(msg: List[BaseModel], meta: List[EventMetadata]):\n",
    "    assert msg == \"Example_msg\"\n",
    "    assert meta == \"Some_meta\"\n",
    "\n",
    "await with_meta(\"Example_msg\", \"Some_meta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fd1a5c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _prepare_callback(\n",
    "    callback: ConsumeCallable\n",
    ") -> AsyncConsumeMeta:\n",
    "    \"\"\"\n",
    "    Prepares a callback to be used in the consumer loop.\n",
    "        1. If callback is sync, asyncify it\n",
    "        2. Wrap the callback into a safe callback for exception handling\n",
    "\n",
    "    Args:\n",
    "        callback: async callable that will be prepared for use in consumer\n",
    "\n",
    "    Returns:\n",
    "        Prepared callback\n",
    "    \"\"\"\n",
    "    async_callback: Union[AsyncConsume, AsyncConsumeMeta] = (\n",
    "        callback if iscoroutinefunction(callback) else asyncer.asyncify(callback)  # type: ignore\n",
    "    )\n",
    "    return _callback_parameters_wrapper(async_callback)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e996f4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if callback is called when wrapped\n",
    "\n",
    "for is_async in [False, True]:\n",
    "    example_msg = \"Example msg\"\n",
    "    callback = AsyncMock() if is_async else Mock()\n",
    "    prepared_callback = _prepare_callback(callback)\n",
    "\n",
    "    with patch(\"__main__.get_type_hints\") as mock:\n",
    "        mock.return_value = {\"msg\": BaseModel}\n",
    "        await prepared_callback(f\"{example_msg}\", \"Some meta\")\n",
    "\n",
    "    callback.assert_called_once_with(f\"{example_msg}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e0977bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def _stream_msgs(  # type: ignore\n",
    "    msgs: Dict[TopicPartition, bytes],\n",
    "    send_stream: anyio.streams.memory.MemoryObjectSendStream[Any],\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Decodes and streams the message and topic to the send_stream.\n",
    "\n",
    "    Args:\n",
    "        msgs:\n",
    "        send_stream:\n",
    "    \"\"\"\n",
    "    for topic_partition, topic_msgs in msgs.items():\n",
    "        topic = topic_partition.topic\n",
    "        try:\n",
    "            await send_stream.send(topic_msgs)\n",
    "        except Exception as e:\n",
    "            logger.warning(\n",
    "                f\"_stream_msgs(): Unexpected exception '{e.__repr__()}' caught and ignored for topic='{topic_partition.topic}', partition='{topic_partition.partition}' and messages: {topic_msgs!r}\"\n",
    "            )\n",
    "\n",
    "\n",
    "def _decode_streamed_msgs(  # type: ignore\n",
    "    msgs: List[ConsumerRecord], msg_type: BaseModel\n",
    ") -> List[BaseModel]:\n",
    "    decoded_msgs = [msg_type.parse_raw(msg.value.decode(\"utf-8\")) for msg in msgs]\n",
    "    return decoded_msgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "335aa93b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sanity check: one msg, one topic\n",
    "\n",
    "with patch(\"anyio.streams.memory.MemoryObjectSendStream.send\") as mock:\n",
    "    send_stream, receive_stream = anyio.create_memory_object_stream()\n",
    "\n",
    "    topic = \"topic_0\"\n",
    "    partition = 0\n",
    "    topic_part_0_0 = TopicPartition(topic, partition)\n",
    "    msg = MyMessage(url=\"http://www.acme.com\", port=22)\n",
    "    record = create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "\n",
    "    await _stream_msgs(\n",
    "        msgs={topic_part_0_0: [record]},\n",
    "        send_stream=send_stream,\n",
    "    )\n",
    "\n",
    "    mock.assert_called_once()\n",
    "    mock.assert_has_calls([call([record])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f25ecc98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check different topics\n",
    "\n",
    "# Two msg, two topics, send called twice with each topic\n",
    "\n",
    "with patch(\"anyio.streams.memory.MemoryObjectSendStream.send\") as mock:\n",
    "    send_stream, receive_stream = anyio.create_memory_object_stream()\n",
    "\n",
    "    topic_partitions = [(\"topic_0\", 0), (\"topic_1\", 0)]\n",
    "\n",
    "    msg = MyMessage(url=\"http://www.acme.com\", port=22)\n",
    "    msgs = {\n",
    "        TopicPartition(topic, partition): [\n",
    "            create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "        ]\n",
    "        for topic, partition in topic_partitions\n",
    "    }\n",
    "\n",
    "    await _stream_msgs(\n",
    "        msgs=msgs,\n",
    "        send_stream=send_stream,\n",
    "    )\n",
    "\n",
    "    assert mock.call_count == 2\n",
    "\n",
    "    mock.assert_has_calls([call(msg) for msg in msgs.values()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff3fa870",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check multiple msgs in same topic\n",
    "\n",
    "# Two msg, one topic, send called twice for same topic\n",
    "\n",
    "with patch(\"anyio.streams.memory.MemoryObjectSendStream.send\") as mock:\n",
    "    send_stream, receive_stream = anyio.create_memory_object_stream()\n",
    "\n",
    "    topic_partitions = [(\"topic_0\", 0)]\n",
    "\n",
    "    msg = MyMessage(url=\"http://www.acme.com\", port=22)\n",
    "    record = create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "\n",
    "    msgs = {\n",
    "        TopicPartition(topic, partition): [\n",
    "            create_consumer_record(topic=topic, partition=partition, msg=msg),\n",
    "            create_consumer_record(topic=topic, partition=partition, msg=msg),\n",
    "        ]\n",
    "        for topic, partition in topic_partitions\n",
    "    }\n",
    "\n",
    "    await _stream_msgs(\n",
    "        msgs=msgs,\n",
    "        send_stream=send_stream,\n",
    "    )\n",
    "\n",
    "    mock.assert_has_calls([call(msg) for msg in msgs.values()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "403988a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check multiple partitions\n",
    "\n",
    "# Two msg, one topic, differenct partitions, send called twice for same topic\n",
    "\n",
    "with patch(\"anyio.streams.memory.MemoryObjectSendStream.send\") as mock:\n",
    "    send_stream, receive_stream = anyio.create_memory_object_stream()\n",
    "\n",
    "    topic_partitions = [(\"topic_0\", 0), (\"topic_0\", 1)]\n",
    "\n",
    "    msg = MyMessage(url=\"http://www.acme.com\", port=22)\n",
    "    msgs = {\n",
    "        TopicPartition(topic, partition): [\n",
    "            create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "        ]\n",
    "        for topic, partition in topic_partitions\n",
    "    }\n",
    "    record = create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "\n",
    "    await _stream_msgs(\n",
    "        msgs=msgs,\n",
    "        send_stream=send_stream,\n",
    "    )\n",
    "\n",
    "    mock.assert_has_calls([call(msg) for msg in msgs.values()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5275963",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _get_single_msg_handlers(  # type: ignore\n",
    "    *,\n",
    "    consumer: fastkafka._aiokafka_imports.AIOKafkaConsumer,\n",
    "    callback: AsyncConsumeMeta,\n",
    "    decoder_fn: Callable[[bytes, Type[BaseModel]], Any],\n",
    "    msg_type: Type[BaseModel],\n",
    "    **kwargs: Any,\n",
    ") -> Tuple[\n",
    "    Callable[\n",
    "        [\n",
    "            ConsumerRecord,\n",
    "            AsyncConsumeMeta,\n",
    "            Callable[[bytes, Type[BaseModel]], Any],\n",
    "            Type[BaseModel],\n",
    "        ],\n",
    "        Awaitable[None],\n",
    "    ],\n",
    "    Callable[\n",
    "        [fastkafka._aiokafka_imports.AIOKafkaConsumer, Any],\n",
    "        Awaitable[List[ConsumerRecord]],\n",
    "    ],\n",
    "]:\n",
    "    \"\"\"\n",
    "    Retrieves the message handlers for consuming single messages from a Kafka topic.\n",
    "\n",
    "    Args:\n",
    "        consumer: The Kafka consumer instance.\n",
    "        callback: The callback function to handle the consumed message.\n",
    "        decoder_fn: The function to decode the consumed message.\n",
    "        msg_type: The type of the consumed message.\n",
    "        **kwargs: Additional keyword arguments for the consumer.\n",
    "\n",
    "    Returns:\n",
    "        The handle_msg function and poll_consumer function.\n",
    "    \"\"\"\n",
    "\n",
    "    async def handle_msg(  # type: ignore\n",
    "        record: ConsumerRecord,\n",
    "        callback: AsyncConsumeMeta = callback,\n",
    "        decoder_fn: Callable[[bytes, Type[BaseModel]], Any] = decoder_fn,\n",
    "        msg_type: Type[BaseModel] = msg_type,\n",
    "    ) -> None:\n",
    "        await callback(\n",
    "            decoder_fn(record.value, msg_type),\n",
    "            EventMetadata.create_event_metadata(record),\n",
    "        )\n",
    "\n",
    "    async def poll_consumer(  # type: ignore\n",
    "        consumer: fastkafka._aiokafka_imports.AIOKafkaConsumer = consumer,\n",
    "        kwargs: Any = kwargs,\n",
    "    ) -> List[ConsumerRecord]:\n",
    "        msgs = await consumer.getmany(**kwargs)\n",
    "        return [msg for msg_group in msgs.values() for msg in msg_group]\n",
    "\n",
    "    return handle_msg, poll_consumer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc43f2f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "topic_partitions = [(\"topic_0\", 0), (\"topic_0\", 1)]\n",
    "\n",
    "msg = MyMessage(url=\"http://www.acme.com\", port=22)\n",
    "msgs = {\n",
    "    TopicPartition(topic, partition): [\n",
    "        create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "    ]\n",
    "    for topic, partition in topic_partitions\n",
    "}\n",
    "record = create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "\n",
    "consumer = AsyncMock()\n",
    "consumer.getmany.return_value = msgs\n",
    "\n",
    "callback = AsyncMock()\n",
    "decoder_fn = json_decoder\n",
    "msg_type = MyMessage\n",
    "\n",
    "handle_msg, poll_consumer = _get_single_msg_handlers(\n",
    "    consumer=consumer, callback=callback, decoder_fn=decoder_fn, msg_type=msg_type\n",
    ")\n",
    "\n",
    "got_msgs = await poll_consumer()\n",
    "assert len(msgs.values()) == len(got_msgs)\n",
    "\n",
    "for msg in got_msgs:\n",
    "    await handle_msg(msg)\n",
    "\n",
    "callback.assert_has_awaits(\n",
    "    [\n",
    "        call(\n",
    "            json_decoder(msg.value, msg_type), EventMetadata.create_event_metadata(msg)\n",
    "        )\n",
    "        for msg in got_msgs\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ddd1b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _get_batch_msg_handlers(  # type: ignore\n",
    "    *,\n",
    "    consumer: fastkafka._aiokafka_imports.AIOKafkaConsumer,\n",
    "    callback: AsyncConsumeMeta,\n",
    "    decoder_fn: Callable[[bytes, Type[BaseModel]], Any],\n",
    "    msg_type: Type[BaseModel],\n",
    "    **kwargs: Any,\n",
    ") -> Tuple[\n",
    "    Callable[\n",
    "        [\n",
    "            List[ConsumerRecord],\n",
    "            AsyncConsumeMeta,\n",
    "            Callable[[bytes, Type[BaseModel]], Any],\n",
    "            Type[BaseModel],\n",
    "        ],\n",
    "        Awaitable[None],\n",
    "    ],\n",
    "    Callable[[fastkafka._aiokafka_imports.AIOKafkaConsumer, Any], Awaitable[List[List[ConsumerRecord]]]],\n",
    "]:\n",
    "    \"\"\"\n",
    "    Retrieves the message handlers for consuming messages in batches from a Kafka topic.\n",
    "\n",
    "    Args:\n",
    "        consumer: The Kafka consumer instance.\n",
    "        callback: The callback function to handle the consumed messages.\n",
    "        decoder_fn: The function to decode the consumed messages.\n",
    "        msg_type: The type of the consumed messages.\n",
    "        **kwargs: Additional keyword arguments for the consumer.\n",
    "\n",
    "    Returns:\n",
    "        The handle_msg function and poll_consumer function.\n",
    "    \"\"\"\n",
    "\n",
    "    async def handle_msg(  # type: ignore\n",
    "        records: List[ConsumerRecord],\n",
    "        callback: AsyncConsumeMeta = callback,\n",
    "        decoder_fn: Callable[[bytes, Type[BaseModel]], Any] = decoder_fn,\n",
    "        msg_type: Type[BaseModel] = msg_type,\n",
    "    ) -> None:\n",
    "        await callback(\n",
    "            [decoder_fn(record.value, msg_type) for record in records],\n",
    "            [EventMetadata.create_event_metadata(record) for record in records],\n",
    "        )\n",
    "\n",
    "    async def poll_consumer(  # type: ignore\n",
    "        consumer: fastkafka._aiokafka_imports.AIOKafkaConsumer = consumer, kwargs: Any = kwargs\n",
    "    ) -> List[List[ConsumerRecord]]:\n",
    "        msgs = await consumer.getmany(**kwargs)\n",
    "        return [value for value in msgs.values() if len(value)>0]\n",
    "\n",
    "    return handle_msg, poll_consumer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18367eb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "topic_partitions = [(\"topic_0\", 0), (\"topic_0\", 1)]\n",
    "\n",
    "msg = MyMessage(url=\"http://www.acme.com\", port=22)\n",
    "msgs = {\n",
    "    TopicPartition(topic, partition): [\n",
    "        create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "    ]\n",
    "    for topic, partition in topic_partitions\n",
    "}\n",
    "record = create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "\n",
    "consumer = AsyncMock()\n",
    "consumer.getmany.return_value = msgs\n",
    "\n",
    "callback = AsyncMock()\n",
    "decoder_fn = json_decoder\n",
    "msg_type = MyMessage\n",
    "\n",
    "handle_msg, poll_consumer = _get_batch_msg_handlers(\n",
    "    consumer=consumer, callback=callback, decoder_fn=decoder_fn, msg_type=msg_type\n",
    ")\n",
    "\n",
    "got_msgs = await poll_consumer()\n",
    "assert len(msgs.values()) == len(got_msgs)\n",
    "\n",
    "for msgs in got_msgs:\n",
    "    assert len(msgs) == 1\n",
    "\n",
    "for msg in got_msgs:\n",
    "    await handle_msg(msg)\n",
    "\n",
    "callback.assert_has_awaits(\n",
    "    [\n",
    "        call(\n",
    "            [json_decoder(msg_unwrapped.value, msg_type) for msg_unwrapped in msg],\n",
    "            [EventMetadata.create_event_metadata(msg_unwrapped) for msg_unwrapped in msg],\n",
    "        )\n",
    "        for msg in got_msgs\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df02ed2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@delegates(fastkafka._aiokafka_imports.AIOKafkaConsumer.getmany)\n",
    "async def _aiokafka_consumer_loop(  # type: ignore\n",
    "    consumer: fastkafka._aiokafka_imports.AIOKafkaConsumer,\n",
    "    *,\n",
    "    topic: str,\n",
    "    decoder_fn: Callable[[bytes, Type[BaseModel]], Any],\n",
    "    callback: ConsumeCallable,\n",
    "    max_buffer_size: int = 100_000,\n",
    "    msg_type: Union[Type[List[BaseModel]], Type[BaseModel]],\n",
    "    is_shutting_down_f: Callable[[], bool],\n",
    "    executor: Union[str, StreamExecutor, None] = None,\n",
    "    **kwargs: Any,\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Consumer loop for infinite pooling of the AIOKafka consumer for new messages. Calls consumer.getmany()\n",
    "    and after the consumer return messages or times out, messages are decoded and streamed to defined callback.\n",
    "\n",
    "    Args:\n",
    "        topic: Topic to subscribe\n",
    "        decoder_fn: Function to decode the messages consumed from the topic\n",
    "        callbacks: Dict of callbacks mapped to their respective topics\n",
    "        timeout_ms: Time to timeut the getmany request by the consumer\n",
    "        max_buffer_size: Maximum number of unconsumed messages in the callback buffer\n",
    "        msg_types: Dict of message types mapped to their respective topics\n",
    "        is_shutting_down_f: Function for controlling the shutdown of consumer loop\n",
    "    \"\"\"\n",
    "\n",
    "    prepared_callback = _prepare_callback(callback)\n",
    "\n",
    "    if hasattr(msg_type, \"__origin__\") and msg_type.__origin__ == list:\n",
    "        handle_msg, poll_consumer = _get_batch_msg_handlers(\n",
    "            consumer=consumer,\n",
    "            callback=prepared_callback,\n",
    "            decoder_fn=decoder_fn,\n",
    "            msg_type=msg_type.__args__[0],  # type: ignore\n",
    "            **kwargs,\n",
    "        )\n",
    "    else:\n",
    "        handle_msg, poll_consumer = _get_single_msg_handlers(\n",
    "            consumer=consumer,\n",
    "            callback=prepared_callback,\n",
    "            decoder_fn=decoder_fn,\n",
    "            msg_type=msg_type,  # type: ignore\n",
    "            **kwargs,\n",
    "        )\n",
    "\n",
    "    await get_executor(executor).run(\n",
    "        is_shutting_down_f=is_shutting_down_f,\n",
    "        generator=poll_consumer,  # type: ignore\n",
    "        processor=handle_msg,  # type: ignore\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4a60f81",
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_shutting_down_f(mock_func: Mock, num_calls: int = 1) -> Callable[[], bool]:\n",
    "    def _is_shutting_down_f():\n",
    "        return mock_func.call_count == num_calls\n",
    "\n",
    "    return _is_shutting_down_f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3020fa4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastkafka._components.task_streaming import SequentialExecutor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77397e6a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ok\n"
     ]
    }
   ],
   "source": [
    "topic = \"topic_0\"\n",
    "partition = 0\n",
    "msg = MyMessage(url=\"http://www.acme.com\", port=22)\n",
    "record = create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "\n",
    "mock_consumer = MagicMock()\n",
    "msgs = {TopicPartition(topic, 0): [record]}\n",
    "\n",
    "f = asyncio.Future()\n",
    "f.set_result(msgs)\n",
    "mock_consumer.configure_mock(**{\"getmany.return_value\": f})\n",
    "\n",
    "def f(msg: MyMessage): pass\n",
    "mock_callback = MagicMock(spec=f)\n",
    "\n",
    "\n",
    "for is_async in [True, False]:\n",
    "    for executor_type in [\"DynamicTaskExecutor\", \"SequentialExecutor\"]:\n",
    "        await _aiokafka_consumer_loop(\n",
    "            consumer=mock_consumer,\n",
    "            topic=topic,\n",
    "            decoder_fn=json_decoder,\n",
    "            max_buffer_size=100,\n",
    "            timeout_ms=10,\n",
    "            callback=asyncer.asyncify(mock_callback) if is_async else mock_callback,\n",
    "            msg_type=MyMessage,\n",
    "            is_shutting_down_f=is_shutting_down_f(mock_consumer.getmany),\n",
    "            executor_type=executor_type,\n",
    "        )\n",
    "\n",
    "        assert mock_consumer.getmany.call_count == 1\n",
    "        mock_callback.assert_called_once_with(msg)\n",
    "\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6854bb7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ok\n"
     ]
    }
   ],
   "source": [
    "topic = \"topic_0\"\n",
    "partition = 0\n",
    "msg = MyMessage(url=\"http://www.acme.com\", port=22)\n",
    "record = create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "\n",
    "mock_consumer = MagicMock()\n",
    "msgs = {TopicPartition(topic, 0): [record]}\n",
    "\n",
    "f = asyncio.Future()\n",
    "f.set_result(msgs)\n",
    "mock_consumer.configure_mock(**{\"getmany.return_value\": f})\n",
    "\n",
    "def f(msg: List[MyMessage]): pass\n",
    "mock_callback = MagicMock(spec=f)\n",
    "\n",
    "\n",
    "for is_async in [True, False]:\n",
    "    for executor_type in [\"DynamicTaskExecutor\", \"SequentialExecutor\"]:\n",
    "        await _aiokafka_consumer_loop(\n",
    "            consumer=mock_consumer,\n",
    "            topic=topic,\n",
    "            decoder_fn=json_decoder,\n",
    "            max_buffer_size=100,\n",
    "            timeout_ms=10,\n",
    "            callback=asyncer.asyncify(mock_callback) if is_async else mock_callback,\n",
    "            msg_type=List[MyMessage],\n",
    "            is_shutting_down_f=is_shutting_down_f(mock_consumer.getmany),\n",
    "            executor_type=executor_type,\n",
    "        )\n",
    "\n",
    "        assert mock_consumer.getmany.call_count == 1\n",
    "        mock_callback.assert_called_once_with([msg])\n",
    "\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54b9e6fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[WARNING] fastkafka._components.task_streaming: e=Exception('')\n",
      "[WARNING] fastkafka._components.task_streaming: e=Exception('')\n",
      "[WARNING] fastkafka._components.task_streaming: e=Exception('')\n",
      "[WARNING] fastkafka._components.task_streaming: e=Exception('')\n",
      "[WARNING] fastkafka._components.task_streaming: e=Exception('')\n",
      "[WARNING] fastkafka._components.task_streaming: e=Exception('')\n",
      "[WARNING] fastkafka._components.task_streaming: e=Exception('')\n",
      "[WARNING] fastkafka._components.task_streaming: e=Exception('')\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "# Sanity check: exception in callback recovery\n",
    "# Two msg, one topic, process_f called twice even tough it throws\n",
    "\n",
    "for is_async in [True, False]:\n",
    "    for executor_type in [\"DynamicTaskExecutor\", \"SequentialExecutor\"]:\n",
    "        topic = \"topic_0\"\n",
    "        partition = 0\n",
    "        msg = MyMessage(url=\"http://www.acme.com\", port=22)\n",
    "        record = create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "\n",
    "        num_msgs = 2\n",
    "\n",
    "        mock_consumer = MagicMock()\n",
    "        msgs = {TopicPartition(topic, 0): [record, record]}\n",
    "\n",
    "        f = asyncio.Future()\n",
    "        f.set_result(msgs)\n",
    "\n",
    "        mock_consumer.configure_mock(**{\"getmany.return_value\": f})\n",
    "        mock_callback = Mock()\n",
    "\n",
    "        exception = Exception(\"\")\n",
    "        mock_callback.side_effect = exception\n",
    "\n",
    "\n",
    "        await _aiokafka_consumer_loop(\n",
    "            consumer=mock_consumer,\n",
    "            topic=topic,\n",
    "            decoder_fn=json_decoder,\n",
    "            max_buffer_size=100,\n",
    "            timeout_ms=1,\n",
    "            callback=asyncer.asyncify(mock_callback) if is_async else mock_callback,\n",
    "            msg_type=MyMessage,\n",
    "            is_shutting_down_f=is_shutting_down_f(mock_consumer.getmany, num_calls=1),\n",
    "            executor_type=executor_type,\n",
    "        )\n",
    "\n",
    "        assert mock_callback.call_count == num_msgs, mock_callback.call_count\n",
    "        mock_callback.assert_has_calls([call(msg), call(msg)])\n",
    "\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2afe654a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ok\n"
     ]
    }
   ],
   "source": [
    "# Sanity check: malformed msgs\n",
    "# One msg of wrong type, two normal msg, one topic, process_f called twice\n",
    "\n",
    "topic = \"topic_0\"\n",
    "partition = 0\n",
    "msg = MyMessage(url=\"http://www.acme.com\", port=22)\n",
    "correct_record = create_consumer_record(topic=topic, partition=partition, msg=msg)\n",
    "faulty_record = create_consumer_record(topic=topic, partition=partition, msg=\"Wrong!\")\n",
    "\n",
    "mock_consumer = MagicMock()\n",
    "msgs = {TopicPartition(topic, 0): [faulty_record, correct_record, correct_record]}\n",
    "\n",
    "mock_consumer.configure_mock(**{\"getmany.return_value\": f})\n",
    "mock_callback = Mock()\n",
    "\n",
    "exception = Exception(\"\")\n",
    "callback.side_effect = exception\n",
    "\n",
    "for is_async in [True, False]:\n",
    "    await _aiokafka_consumer_loop(\n",
    "        consumer=mock_consumer,\n",
    "        topic=topic,\n",
    "        decoder_fn=json_decoder,\n",
    "        max_buffer_size=100,\n",
    "        timeout_ms=10,\n",
    "        callback=asyncer.asyncify(mock_callback) if is_async else mock_callback,\n",
    "        msg_type=MyMessage,\n",
    "        is_shutting_down_f=is_shutting_down_f(mock_consumer.getmany),\n",
    "    )\n",
    "\n",
    "    assert mock_consumer.getmany.call_count == 1\n",
    "    mock_callback.assert_has_calls([call(msg), call(msg)])\n",
    "\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46031397",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def sanitize_kafka_config(**kwargs: Any) -> Dict[str, Any]:\n",
    "    \"\"\"Sanitize Kafka config\"\"\"\n",
    "    return {k: \"*\" * len(v) if \"pass\" in k.lower() else v for k, v in kwargs.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfa2ec97",
   "metadata": {},
   "outputs": [],
   "source": [
    "kwargs = {\n",
    "    \"bootstrap_servers\": \"whatever.cloud:9092\",\n",
    "    \"auto_offset_reset\": \"earliest\",\n",
    "    \"security_protocol\": \"SASL_SSL\",\n",
    "    \"sasl_mechanism\": \"PLAIN\",\n",
    "    \"sasl_plain_username\": \"username\",\n",
    "    \"sasl_plain_password\": \"password\",\n",
    "    \"ssl_context\": \"something\",\n",
    "}\n",
    "\n",
    "assert sanitize_kafka_config(**kwargs)[\"sasl_plain_password\"] == \"********\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca7ba3a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@delegates(fastkafka._aiokafka_imports.AIOKafkaConsumer)\n",
    "@delegates(_aiokafka_consumer_loop, keep=True)\n",
    "async def aiokafka_consumer_loop(\n",
    "    topic: str,\n",
    "    decoder_fn: Callable[[bytes, Type[BaseModel]], Any],\n",
    "    *,\n",
    "    timeout_ms: int = 100,\n",
    "    max_buffer_size: int = 100_000,\n",
    "    callback: ConsumeCallable,\n",
    "    msg_type: Union[Type[List[BaseModel]], Type[BaseModel]],\n",
    "    is_shutting_down_f: Callable[[], bool],\n",
    "    executor: Union[str, StreamExecutor, None] = None,\n",
    "    **kwargs: Any,\n",
    ") -> None:\n",
    "    \"\"\"Consumer loop for infinite pooling of the AIOKafka consumer for new messages. Creates and starts AIOKafkaConsumer\n",
    "    and runs _aio_kafka_consumer loop fo infinite poling of the consumer for new messages.\n",
    "\n",
    "    Args:\n",
    "        topic: name of the topic to subscribe to\n",
    "        decoder_fn: Function to decode the messages consumed from the topic\n",
    "        callback: callback function to be called after decoding and parsing a consumed message\n",
    "        timeout_ms: Time to timeut the getmany request by the consumer\n",
    "        max_buffer_size: Maximum number of unconsumed messages in the callback buffer\n",
    "        msg_type: Type with `parse_json` method used for parsing a decoded message\n",
    "        is_shutting_down_f: Function for controlling the shutdown of consumer loop\n",
    "    \"\"\"\n",
    "    logger.info(f\"aiokafka_consumer_loop() starting...\")\n",
    "    try:\n",
    "        consumer = fastkafka._aiokafka_imports.AIOKafkaConsumer(\n",
    "            **kwargs,\n",
    "        )\n",
    "        logger.info(\n",
    "            f\"aiokafka_consumer_loop(): Consumer created using the following parameters: {sanitize_kafka_config(**kwargs)}\"\n",
    "        )\n",
    "        \n",
    "        await consumer.start()\n",
    "        logger.info(\"aiokafka_consumer_loop(): Consumer started.\")\n",
    "        consumer.subscribe([topic])\n",
    "        logger.info(\"aiokafka_consumer_loop(): Consumer subscribed.\")\n",
    "\n",
    "        try:\n",
    "            await _aiokafka_consumer_loop(\n",
    "                consumer=consumer,\n",
    "                topic=topic,\n",
    "                decoder_fn=decoder_fn,\n",
    "                max_buffer_size=max_buffer_size,\n",
    "                timeout_ms=timeout_ms,\n",
    "                callback=callback,\n",
    "                msg_type=msg_type,\n",
    "                is_shutting_down_f=is_shutting_down_f,\n",
    "                executor = executor,\n",
    "            )\n",
    "        finally:\n",
    "            await consumer.stop()\n",
    "            logger.info(f\"aiokafka_consumer_loop(): Consumer stopped.\")\n",
    "            logger.info(f\"aiokafka_consumer_loop() finished.\")\n",
    "    except Exception as e:\n",
    "        logger.error(\n",
    "            f\"aiokafka_consumer_loop(): unexpected exception raised: '{e.__repr__()}'\"\n",
    "        )\n",
    "        raise e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc89d47e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: But not exported to PATH, exporting...\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._components.test_dependencies: But not exported to PATH, exporting...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b0f0bbb3ee124e0c8d59975fd0b37656",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/9178 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: msgs_received=1000\n",
      "[INFO] __main__: msgs_received=2000\n",
      "[INFO] __main__: msgs_received=3000\n",
      "[INFO] __main__: msgs_received=4000\n",
      "[INFO] __main__: msgs_received=5000\n",
      "[INFO] __main__: msgs_received=6000\n",
      "[INFO] __main__: msgs_received=7000\n",
      "[INFO] __main__: msgs_received=8000\n",
      "[INFO] __main__: msgs_received=9000\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 228127...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 228127 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 227746...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 227746 terminated.\n",
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb5162a2670f4f93a5c61b5e8d5de0b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/9178 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: msgs_received=1000\n",
      "[INFO] __main__: msgs_received=2000\n",
      "[INFO] __main__: msgs_received=3000\n",
      "[INFO] __main__: msgs_received=4000\n",
      "[INFO] __main__: msgs_received=5000\n",
      "[INFO] __main__: msgs_received=6000\n",
      "[INFO] __main__: msgs_received=7000\n",
      "[INFO] __main__: msgs_received=8000\n",
      "[INFO] __main__: msgs_received=9000\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 229358...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 229358 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 228977...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 228977 terminated.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "for executor in [\"DynamicTaskExecutor\", \"SequentialExecutor\"]:\n",
    "    topic = \"test_topic\"\n",
    "    msgs_sent = 9178\n",
    "    msgs = [\n",
    "        MyMessage(url=\"http://www.ai.com\", port=port).model_dump_json().encode(\"utf-8\")\n",
    "        for port in range(msgs_sent)\n",
    "    ]\n",
    "    msgs_received = 0\n",
    "\n",
    "\n",
    "    async def count_msg(msg: MyMessage):\n",
    "        global msgs_received\n",
    "        msgs_received = msgs_received + 1\n",
    "        if msgs_received % 1000 == 0:\n",
    "            logger.info(f\"{msgs_received=}\")\n",
    "\n",
    "\n",
    "    async with ApacheKafkaBroker(topics=[topic], listener_port=11992) as bootstrap_server:\n",
    "        await produce_messages(topic=topic, bootstrap_servers=bootstrap_server, msgs=msgs)\n",
    "        await aiokafka_consumer_loop(\n",
    "            topic=topic,\n",
    "            decoder_fn=json_decoder,\n",
    "            auto_offset_reset=\"earliest\",\n",
    "            callback=count_msg,\n",
    "            msg_type=MyMessage,\n",
    "            is_shutting_down_f=true_after(2),\n",
    "            bootstrap_servers=bootstrap_server,\n",
    "            executor=executor,\n",
    "        )\n",
    "\n",
    "        assert msgs_sent == msgs_received, f\"{msgs_sent} != {msgs_received}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fde91e3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cf8d3ff86a5c4328b69a58fa64bf008b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/9178 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: msgs_received=9178\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 230586...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 230586 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 230206...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 230206 terminated.\n",
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f2bfe58038548168f841826c32b63ba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/9178 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: msgs_received=9178\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 231816...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 231816 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 231435...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 231435 terminated.\n"
     ]
    }
   ],
   "source": [
    "for executor in [\"DynamicTaskExecutor\", \"SequentialExecutor\"]:\n",
    "    topic = \"test_topic\"\n",
    "    msgs_sent = 9178\n",
    "    msgs = [\n",
    "        MyMessage(url=\"http://www.ai.com\", port=port).model_dump_json().encode(\"utf-8\")\n",
    "        for port in range(msgs_sent)\n",
    "    ]\n",
    "    msgs_received = 0\n",
    "\n",
    "\n",
    "    async def count_msg(msg: List[MyMessage], meta: List[EventMetadata]):\n",
    "        global msgs_received\n",
    "        msgs_received = msgs_received + len(msg)\n",
    "        logger.info(f\"{msgs_received=}\")\n",
    "\n",
    "\n",
    "    async with ApacheKafkaBroker(topics=[topic], listener_port=11992) as bootstrap_server:\n",
    "        await produce_messages(topic=topic, bootstrap_servers=bootstrap_server, msgs=msgs)\n",
    "        await aiokafka_consumer_loop(\n",
    "            topic=topic,\n",
    "            decoder_fn=json_decoder,\n",
    "            auto_offset_reset=\"earliest\",\n",
    "            callback=count_msg,\n",
    "            msg_type=List[MyMessage],\n",
    "            is_shutting_down_f=true_after(2),\n",
    "            bootstrap_servers=bootstrap_server,\n",
    "            executor=executor,\n",
    "        )\n",
    "\n",
    "        assert msgs_sent == msgs_received, f\"{msgs_sent} != {msgs_received}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e47e37ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ba8321cf4ef4461b8042a3d79ce3d22",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/9178 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: msgs_received=1000, meta=EventMetadata(topic='test_topic', partition=0, offset=999, timestamp=1688404081071, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":999}', checksum=None, serialized_key_size=-1, serialized_value_size=39, headers=())\n",
      "[INFO] __main__: msgs_received=2000, meta=EventMetadata(topic='test_topic', partition=0, offset=1999, timestamp=1688404081080, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":1999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=3000, meta=EventMetadata(topic='test_topic', partition=0, offset=2999, timestamp=1688404081087, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":2999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=4000, meta=EventMetadata(topic='test_topic', partition=0, offset=3999, timestamp=1688404081095, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":3999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=5000, meta=EventMetadata(topic='test_topic', partition=0, offset=4999, timestamp=1688404081102, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":4999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=6000, meta=EventMetadata(topic='test_topic', partition=0, offset=5999, timestamp=1688404081110, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":5999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=7000, meta=EventMetadata(topic='test_topic', partition=0, offset=6999, timestamp=1688404081116, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":6999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=8000, meta=EventMetadata(topic='test_topic', partition=0, offset=7999, timestamp=1688404081123, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":7999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=9000, meta=EventMetadata(topic='test_topic', partition=0, offset=8999, timestamp=1688404081130, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":8999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 233044...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 233044 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 232663...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 232663 terminated.\n",
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "727b139b07e244af90546172944ad4db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/9178 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: msgs_received=1000, meta=EventMetadata(topic='test_topic', partition=0, offset=999, timestamp=1688404089840, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":999}', checksum=None, serialized_key_size=-1, serialized_value_size=39, headers=())\n",
      "[INFO] __main__: msgs_received=2000, meta=EventMetadata(topic='test_topic', partition=0, offset=1999, timestamp=1688404089848, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":1999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=3000, meta=EventMetadata(topic='test_topic', partition=0, offset=2999, timestamp=1688404089854, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":2999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=4000, meta=EventMetadata(topic='test_topic', partition=0, offset=3999, timestamp=1688404089860, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":3999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=5000, meta=EventMetadata(topic='test_topic', partition=0, offset=4999, timestamp=1688404089867, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":4999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=6000, meta=EventMetadata(topic='test_topic', partition=0, offset=5999, timestamp=1688404089873, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":5999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=7000, meta=EventMetadata(topic='test_topic', partition=0, offset=6999, timestamp=1688404089879, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":6999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=8000, meta=EventMetadata(topic='test_topic', partition=0, offset=7999, timestamp=1688404089886, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":7999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: msgs_received=9000, meta=EventMetadata(topic='test_topic', partition=0, offset=8999, timestamp=1688404089892, timestamp_type=0, key=None, value=b'{\"url\":\"http://www.ai.com/\",\"port\":8999}', checksum=None, serialized_key_size=-1, serialized_value_size=40, headers=())\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 234272...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 234272 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 233892...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 233892 terminated.\n"
     ]
    }
   ],
   "source": [
    "# Test with meta\n",
    "\n",
    "for executor in [\"DynamicTaskExecutor\", \"SequentialExecutor\"]:\n",
    "    topic = \"test_topic\"\n",
    "    msgs_sent = 9178\n",
    "    msgs = [\n",
    "        MyMessage(url=\"http://www.ai.com\", port=port).model_dump_json().encode(\"utf-8\")\n",
    "        for port in range(msgs_sent)\n",
    "    ]\n",
    "    msgs_received = 0\n",
    "    meta_samples = []\n",
    "\n",
    "    async def count_msg(msg: MyMessage, meta: EventMetadata):\n",
    "        global msgs_received\n",
    "        msgs_received = msgs_received + 1\n",
    "        if msgs_received % 1000 == 0:\n",
    "            meta_samples.append(meta)\n",
    "            logger.info(f\"{msgs_received=}, {meta=}\")\n",
    "\n",
    "\n",
    "    async with ApacheKafkaBroker(topics=[topic], listener_port=11992) as bootstrap_server:\n",
    "        await produce_messages(topic=topic, bootstrap_servers=bootstrap_server, msgs=msgs)\n",
    "        await aiokafka_consumer_loop(\n",
    "            topic=topic,\n",
    "            decoder_fn=json_decoder,\n",
    "            auto_offset_reset=\"earliest\",\n",
    "            callback=count_msg,\n",
    "            msg_type=MyMessage,\n",
    "            is_shutting_down_f=true_after(2),\n",
    "            bootstrap_servers=bootstrap_server,\n",
    "            executor = executor\n",
    "        )\n",
    "\n",
    "        assert msgs_sent == msgs_received, f\"{msgs_sent} != {msgs_received}\"\n",
    "        assert all(isinstance(meta, EventMetadata) for meta in meta_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6484e52",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "18826d9659c54c669419983850303cc8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/9178 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: msgs_received=1000\n",
      "[INFO] __main__: msgs_received=2000\n",
      "[INFO] __main__: msgs_received=3000\n",
      "[INFO] __main__: msgs_received=4000\n",
      "[INFO] __main__: msgs_received=5000\n",
      "[INFO] __main__: msgs_received=6000\n",
      "[INFO] __main__: msgs_received=7000\n",
      "[INFO] __main__: msgs_received=8000\n",
      "[INFO] __main__: msgs_received=9000\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 235506...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 235506 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 235124...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 235124 terminated.\n",
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "29a73a0af3b54d858a814c1f0d861d3d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/9178 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: msgs_received=1000\n",
      "[INFO] __main__: msgs_received=2000\n",
      "[INFO] __main__: msgs_received=3000\n",
      "[INFO] __main__: msgs_received=4000\n",
      "[INFO] __main__: msgs_received=5000\n",
      "[INFO] __main__: msgs_received=6000\n",
      "[INFO] __main__: msgs_received=7000\n",
      "[INFO] __main__: msgs_received=8000\n",
      "[INFO] __main__: msgs_received=9000\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 236738...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 236738 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 236357...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 236357 terminated.\n"
     ]
    }
   ],
   "source": [
    "# Test with avro_decoder\n",
    "\n",
    "for executor in [\"DynamicTaskExecutor\", \"SequentialExecutor\"]:\n",
    "    topic = \"test_topic\"\n",
    "    msgs_sent = 9178\n",
    "    msgs = [\n",
    "        avro_encoder(MyMessage(url=\"http://www.ai.com\", port=port))\n",
    "        for port in range(msgs_sent)\n",
    "    ]\n",
    "    msgs_received = 0\n",
    "\n",
    "\n",
    "    async def count_msg(msg: MyMessage):\n",
    "        global msgs_received\n",
    "        msgs_received = msgs_received + 1\n",
    "        if msgs_received % 1000 == 0:\n",
    "            logger.info(f\"{msgs_received=}\")\n",
    "\n",
    "\n",
    "    async with ApacheKafkaBroker(topics=[topic], listener_port=11992) as bootstrap_server:\n",
    "        await produce_messages(topic=topic, bootstrap_servers=bootstrap_server, msgs=msgs)\n",
    "        await aiokafka_consumer_loop(\n",
    "            topic=topic,\n",
    "            decoder_fn=avro_decoder,\n",
    "            auto_offset_reset=\"earliest\",\n",
    "            callback=count_msg,\n",
    "            msg_type=MyMessage,\n",
    "            is_shutting_down_f=true_after(2),\n",
    "            bootstrap_servers=bootstrap_server,\n",
    "            executor=executor,\n",
    "        )\n",
    "\n",
    "        assert msgs_sent == msgs_received, f\"{msgs_sent} != {msgs_received}\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94ea86a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cd276ce4985d4bd3a83130625e28a9cd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/9178 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: msgs_received=1000\n",
      "[INFO] __main__: msgs_received=2000\n",
      "[INFO] __main__: msgs_received=3000\n",
      "[INFO] __main__: msgs_received=4000\n",
      "[INFO] __main__: msgs_received=5000\n",
      "[INFO] __main__: msgs_received=6000\n",
      "[INFO] __main__: msgs_received=7000\n",
      "[INFO] __main__: msgs_received=8000\n",
      "[INFO] __main__: msgs_received=9000\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 237967...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 237967 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 237585...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 237585 terminated.\n",
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4937733125b945ad9dace85f5b570516",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/9178 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: msgs_received=1000\n",
      "[INFO] __main__: msgs_received=2000\n",
      "[INFO] __main__: msgs_received=3000\n",
      "[INFO] __main__: msgs_received=4000\n",
      "[INFO] __main__: msgs_received=5000\n",
      "[INFO] __main__: msgs_received=6000\n",
      "[INFO] __main__: msgs_received=7000\n",
      "[INFO] __main__: msgs_received=8000\n",
      "[INFO] __main__: msgs_received=9000\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 239194...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 239194 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 238813...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 238813 terminated.\n"
     ]
    }
   ],
   "source": [
    "# Test with avro_decoder and meta\n",
    "\n",
    "for executor in [\"DynamicTaskExecutor\", \"SequentialExecutor\"]:\n",
    "    topic = \"test_topic\"\n",
    "    msgs_sent = 9178\n",
    "    msgs = [\n",
    "        avro_encoder(MyMessage(url=\"http://www.ai.com\", port=port))\n",
    "        for port in range(msgs_sent)\n",
    "    ]\n",
    "    msgs_received = 0\n",
    "    meta_samples = []\n",
    "\n",
    "    async def count_msg(msg: MyMessage, meta: EventMetadata):\n",
    "        global msgs_received\n",
    "        msgs_received = msgs_received + 1\n",
    "        if msgs_received % 1000 == 0:\n",
    "            logger.info(f\"{msgs_received=}\")\n",
    "            meta_samples.append(meta)\n",
    "\n",
    "\n",
    "    async with ApacheKafkaBroker(topics=[topic], listener_port=11992) as bootstrap_server:\n",
    "        await produce_messages(topic=topic, bootstrap_servers=bootstrap_server, msgs=msgs)\n",
    "        await aiokafka_consumer_loop(\n",
    "            topic=topic,\n",
    "            decoder_fn=avro_decoder,\n",
    "            auto_offset_reset=\"earliest\",\n",
    "            callback=count_msg,\n",
    "            msg_type=MyMessage,\n",
    "            is_shutting_down_f=true_after(2),\n",
    "            bootstrap_servers=bootstrap_server,\n",
    "            executor=executor,\n",
    "        )\n",
    "\n",
    "        assert msgs_sent == msgs_received, f\"{msgs_sent} != {msgs_received}\"\n",
    "        assert all(isinstance(meta, EventMetadata) for meta in meta_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cd9d9d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02bc2b22db4a4832aafe07b9022a74f4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/50000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d41ef425a1194d138ca92b02c0a4cad0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "consuming messages:   0%|          | 0/50000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "Messages processed: 50,000\n",
      "Time              : 2.36 s\n",
      "Throughput.       : 21,212 msg/s\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 240422...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 240422 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 240041...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 240041 terminated.\n",
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb2981e9f67047bc8df96e629cc0c505",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/50000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "de913723000e49adbe63c93ef447f1a7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "consuming messages:   0%|          | 0/50000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "Messages processed: 50,000\n",
      "Time              : 1.65 s\n",
      "Throughput.       : 30,371 msg/s\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 241652...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 241652 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 241271...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 241271 terminated.\n"
     ]
    }
   ],
   "source": [
    "for executor in [\"DynamicTaskExecutor\", \"SequentialExecutor\"]:\n",
    "    topic = \"test_topic\"\n",
    "    msgs_sent = 500_00\n",
    "    msgs = [\n",
    "        MyMessage(url=\"http://www.ai.com\", port=port).model_dump_json().encode(\"utf-8\")\n",
    "        for port in range(msgs_sent)\n",
    "    ]\n",
    "\n",
    "\n",
    "    async def count_msg(msg: MyMessage):\n",
    "        pbar.update(1)\n",
    "\n",
    "\n",
    "    def _is_shutting_down_f():\n",
    "        return pbar.n >= pbar.total\n",
    "\n",
    "\n",
    "    async with ApacheKafkaBroker(topics=[topic], listener_port=11992) as bootstrap_server:\n",
    "        await produce_messages(topic=topic, bootstrap_servers=bootstrap_server, msgs=msgs)\n",
    "        with tqdm(total=msgs_sent, desc=\"consuming messages\") as _pbar:\n",
    "            global pbar\n",
    "            pbar = _pbar\n",
    "\n",
    "            start = datetime.now()\n",
    "            await aiokafka_consumer_loop(\n",
    "                topic=topic,\n",
    "                decoder_fn=json_decoder,\n",
    "                auto_offset_reset=\"earliest\",\n",
    "                callback=count_msg,\n",
    "                msg_type=MyMessage,\n",
    "                is_shutting_down_f=_is_shutting_down_f,\n",
    "                bootstrap_servers=bootstrap_server,\n",
    "                executor=executor\n",
    "            )\n",
    "            t = (datetime.now() - start) / timedelta(seconds=1)\n",
    "            thrp = pbar.n / t\n",
    "\n",
    "            print(f\"Messages processed: {pbar.n:,d}\")\n",
    "            print(f\"Time              : {t:.2f} s\")\n",
    "            print(f\"Throughput.       : {thrp:,.0f} msg/s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff5ada17",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1dcc60d607ca40d282062e5346709535",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/50000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3daf4e7d24d54342836d35c637035c53",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "consuming messages:   0%|          | 0/50000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "Messages processed: 50,000\n",
      "Time              : 9.75 s\n",
      "Throughput.       : 5,127 msg/s\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 242885...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 242885 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 242503...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 242503 terminated.\n",
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6180748bfe2c46559df2d40406b90a2a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/50000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8e232f75af3d44789a93157e5e88b974",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "consuming messages:   0%|          | 0/50000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "Messages processed: 50,000\n",
      "Time              : 8.43 s\n",
      "Throughput.       : 5,934 msg/s\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 244120...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 244120 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 243738...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 243738 terminated.\n"
     ]
    }
   ],
   "source": [
    "# Test with avro_decoder\n",
    "\n",
    "for executor in [\"DynamicTaskExecutor\", \"SequentialExecutor\"]:\n",
    "    topic = \"test_topic\"\n",
    "    msgs_sent = 500_00\n",
    "    msgs = [\n",
    "        avro_encoder(MyMessage(url=\"http://www.ai.com\", port=port))\n",
    "        for port in range(msgs_sent)\n",
    "    ]\n",
    "\n",
    "\n",
    "    async def count_msg(msg: MyMessage):\n",
    "        pbar.update(1)\n",
    "\n",
    "\n",
    "    def _is_shutting_down_f():\n",
    "        return pbar.n >= pbar.total\n",
    "\n",
    "\n",
    "    async with ApacheKafkaBroker(topics=[topic], listener_port=11992) as bootstrap_server:\n",
    "        await produce_messages(topic=topic, bootstrap_servers=bootstrap_server, msgs=msgs)\n",
    "        with tqdm(total=msgs_sent, desc=\"consuming messages\") as _pbar:\n",
    "            global pbar\n",
    "            pbar = _pbar\n",
    "\n",
    "            start = datetime.now()\n",
    "            await aiokafka_consumer_loop(\n",
    "                topic=topic,\n",
    "                decoder_fn=avro_decoder,\n",
    "                auto_offset_reset=\"earliest\",\n",
    "                callback=count_msg,\n",
    "                msg_type=MyMessage,\n",
    "                is_shutting_down_f=_is_shutting_down_f,\n",
    "                bootstrap_servers=bootstrap_server,\n",
    "                executor=executor\n",
    "            )\n",
    "            t = (datetime.now() - start) / timedelta(seconds=1)\n",
    "            thrp = pbar.n / t\n",
    "\n",
    "            print(f\"Messages processed: {pbar.n:,d}\")\n",
    "            print(f\"Time              : {t:.2f} s\")\n",
    "            print(f\"Throughput.       : {thrp:,.0f} msg/s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac591139",
   "metadata": {},
   "source": [
    "## Consumer loop benchmark and coroutine sanity check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "810bb69a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:11992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "38effce230a74bffa5882effc66009f2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "producing to 'test_topic':   0%|          | 0/50000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "69f36e173de347e88287113ed705b33b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "consuming messages:   0%|          | 0/100000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: aiokafka_consumer_loop() starting...\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'bootstrap_servers': '127.0.0.1:11992'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'test_topic'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'test_topic'}\n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'test_topic': 1}. \n",
      "[INFO] __main__: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] __main__: aiokafka_consumer_loop() finished.\n",
      "Messages processed: 100,000\n",
      "Time              : 5.07 s\n",
      "Throughput.       : 19,721 msg/s\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 245352...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 245352 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 244970...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 244970 terminated.\n"
     ]
    }
   ],
   "source": [
    "# | notest\n",
    "\n",
    "topic = \"test_topic\"\n",
    "msgs_sent = 500_00\n",
    "msgs = [\n",
    "    MyMessage(url=\"http://www.ai.com\", port=port).model_dump_json().encode(\"utf-8\")\n",
    "    for port in range(msgs_sent)\n",
    "]\n",
    "\n",
    "\n",
    "async def count_msg(msg: MyMessage):\n",
    "    pbar.update(1)\n",
    "    await asyncio.sleep(1)\n",
    "    pbar.update(1)\n",
    "\n",
    "def _is_shutting_down_f():\n",
    "    return pbar.n >= pbar.total\n",
    "\n",
    "\n",
    "async with ApacheKafkaBroker(topics=[topic], listener_port=11992) as bootstrap_server:\n",
    "    await produce_messages(topic=topic, bootstrap_servers=bootstrap_server, msgs=msgs)\n",
    "    with tqdm(total=msgs_sent*2, desc=\"consuming messages\") as _pbar:\n",
    "        global pbar\n",
    "        pbar = _pbar\n",
    "\n",
    "        start = datetime.now()\n",
    "        await aiokafka_consumer_loop(\n",
    "            topic=topic,\n",
    "            decoder_fn=json_decoder,\n",
    "            auto_offset_reset=\"earliest\",\n",
    "            callback=count_msg,\n",
    "            msg_type=MyMessage,\n",
    "            is_shutting_down_f=_is_shutting_down_f,\n",
    "            bootstrap_servers=bootstrap_server,\n",
    "            executor = \"DynamicTaskExecutor\"\n",
    "        )\n",
    "        t = (datetime.now() - start) / timedelta(seconds=1)\n",
    "        thrp = pbar.n / t\n",
    "\n",
    "        print(f\"Messages processed: {pbar.n:,d}\")\n",
    "        print(f\"Time              : {t:.2f} s\")\n",
    "        print(f\"Throughput.       : {thrp:,.0f} msg/s\")\n",
    "        \n",
    "assert t < 15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9067015",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
